package cn.bugstack.xfg.dev.tech.config;

import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.OllamaEmbeddingClient;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.PgVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

/**
 * 配置Ollama AI服务所需的各种组件
 */
@Configuration
public class OllamaConfig {

    /**
     * 创建Ollama API客户端，连接到指定的Ollama服务地址
     *
     * @param baseUrl 从配置文件读取的Ollama服务基础URL
     * @return 配置好的OllamaApi实例
     */
    @Bean
    public OllamaApi ollamaApi(@Value("${spring.ai.ollama.base-url}") String baseUrl) {
        return new OllamaApi(baseUrl);
    }

    /**
     * 创建Ollama聊天客户端，用于与Ollama模型进行交互
     *
     * @param ollamaApi 已配置的Ollama API客户端
     * @return 配置好的OllamaChatClient实例
     */
    @Bean
    public OllamaChatClient ollamaChatClient(OllamaApi ollamaApi) {
        return new OllamaChatClient(ollamaApi);
    }

    /**
     * 创建文本分割器，用于将文本按token进行拆分
     * 常用于处理需要分段处理的大文本
     *
     * @return TokenTextSplitter实例
     */
    @Bean
    public TokenTextSplitter tokenTextSplitter() {
        return new TokenTextSplitter();
    }

    /**
     * 创建简单的内存向量存储，使用Ollama模型生成文本嵌入向量
     * 适用于小型应用或开发测试环境
     *
     * @param ollamaApi 已配置的Ollama API客户端
     * @return 配置好的SimpleVectorStore实例
     */
    @Bean
    public SimpleVectorStore simpleVectorStore(OllamaApi ollamaApi) {
        // 配置Ollama嵌入客户端，使用nomic-embed-text模型生成向量
        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
        embeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
        return new SimpleVectorStore(embeddingClient);
    }

    /**
     * 创建基于PostgreSQL的向量存储，使用Ollama模型生成文本嵌入向量
     * 适用于需要持久化存储和高性能查询的生产环境
     *
     * @param ollamaApi    已配置的Ollama API客户端
     * @param jdbcTemplate Spring JDBC模板，用于与PostgreSQL数据库交互
     * @return 配置好的PgVectorStore实例
     */
    @Bean
    public PgVectorStore pgVectorStore(OllamaApi ollamaApi, JdbcTemplate jdbcTemplate) {
        // 配置Ollama嵌入客户端，使用nomic-embed-text模型生成向量
        OllamaEmbeddingClient embeddingClient = new OllamaEmbeddingClient(ollamaApi);
        embeddingClient.withDefaultOptions(OllamaOptions.create().withModel("nomic-embed-text"));
        return new PgVectorStore(jdbcTemplate, embeddingClient);
    }
}
